package middle;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.PriorityQueue;
import java.util.Set;

public class Solution_355 {

    public static void main(String[] args) {
        Twitter model =  new Twitter();
//        model.postTweet(1, 5);
//        model.postTweet(1, 3);
//        model.postTweet(1, 101);
//        model.postTweet(1, 13);
//        model.postTweet(1, 10);
//        model.postTweet(1, 2);
//        model.postTweet(1, 94);
//        model.postTweet(1, 505);
//        model.postTweet(1, 333);
//        model.postTweet(1, 22);
//        model.postTweet(1, 11);
//        System.out.println(model.getNewsFeed(1));
        model.postTweet(1, 5);
        System.out.println(model.getNewsFeed(1));
        model.follow(1, 2);
        model.postTweet(2, 6);
        System.out.println(model.getNewsFeed(1));
        model.unfollow(1, 2);
        System.out.println(model.getNewsFeed(1));
    }

    static class Twitter {
        private int index = 1;
        private final Map<Integer, Set<Integer>> followerMap;
        private final Map<Integer, PriorityQueue<int[]>> tweetMap;

        public Twitter() {
            followerMap = new HashMap<>();
            tweetMap = new HashMap<>();
        }

        public void postTweet(int userId, int tweetId) {
            PriorityQueue<int[]> queue = tweetMap.computeIfAbsent(userId, k -> new PriorityQueue<>(
                Comparator.comparingInt(a -> a[1])));
            queue.add(new int[]{tweetId, index++});
            if (queue.size() > 10) {
                queue.poll();
            }
        }

        public List<Integer> getNewsFeed(int userId) {
            PriorityQueue<int[]> queue = new PriorityQueue<>(Comparator.comparingInt(a -> a[1]));
            add(queue, userId);
            Set<Integer> set = followerMap.get(userId);
            if (set != null) {
                for (Integer id : set) {
                    add(queue, id);
                }
            }
            List<Integer> list = new ArrayList<>(queue.size());
            while (!queue.isEmpty()) {
                list.add(0, queue.poll()[0]);
            }
            return list;
        }

        private void add(PriorityQueue<int[]> queue, int userId) {
            PriorityQueue<int[]> tweetQueue = tweetMap.get(userId);
            if (tweetQueue != null) {
                for (int[] ints : tweetQueue) {
                    queue.add(ints);
                    if (queue.size() > 10) {
                        queue.poll();
                    }
                }
            }
        }

        public void follow(int followerId, int followeeId) {
            Set<Integer> set = followerMap.computeIfAbsent(followerId, k -> new HashSet<>());
            set.add(followeeId);
        }

        public void unfollow(int followerId, int followeeId) {
            Set<Integer> set = followerMap.get(followerId);
            if (set != null) {
                set.remove(followeeId);
            }
        }
    }
}
